from flask import Flask, request, jsonify, render_template_string
from Unit.embedding import encode_code_and_save_to_csv
from Unit.searchRAG import search_milvus
from Unit.checkCode import check_code
from Unit.reflection import reflection_code
from Unit.cot import get_cot
from Fuzz.testDriver import get_driver_data
from Fuzz.execute_code import execute_code
import time
import os



# if os.path.exists("test_results.txt"):
#   os.remove("test_results.txt")

app = Flask(__name__)

def id_transform_file(id):
    padded_string = str(id).rjust(5, '0')
    return f"data/testCase_{padded_string}.txt"

# 接收前端收到的文件
@app.route('/process-code', methods=['POST'])
def process_code():
    if 'codeFile' not in request.files:
        return "No file part", 400
    file = request.files['codeFile']
    if file.filename == '':
        return "No selected file", 400
    if file:

        # 读取文件
        code_snippet = file.read().decode('utf-8')

        # 编码处理
        output_csv = "function_vector.csv"  # 输出的 CSV 文件
        embedding = encode_code_and_save_to_csv(code_snippet, output_csv)

        # Milvus 搜索
        # TODO: 连接Milvus向量数据库
        search_id = search_milvus(embedding)

        # 检索到相似功能函数
        target_file = id_transform_file(search_id)

        # CoT : 生成测试用例
        tested_func = get_cot(code_snippet, target_file)
        print("tested_func:")
        print(tested_func)
        # 执行测试用例
        # 测试驱动代码
        driver_code = get_driver_data(tested_func, code_snippet)
        print("driver_code:")
        print(driver_code)
        # # reflected_code = reflection_code(driver_code[0])
        #
        #checked_code = check_code(reflected_code)
        # 大模型检查代码正确性
        checked_code = check_code(driver_code)
        print("checked_code:")
        print(checked_code)
        # # 反射，怎么拿到反射后的数据结构？
        # # 反射，怎么拿到反射后的数据结构？
        # # 自定义，库类型，怎么处理？
        #
        # # 执行驱动函数生成变异数据
        # # 初始种子
        execute_code(checked_code)
        # # execute_code(reflected_code)

            # 传统模糊测试工具
            # TODO : 将得到的变异数据传给传统模糊测试工具

        return jsonify({"message": "Code processed successfully.", "embedding": embedding})
    else:
        return "Invalid file type", 400

@app.route('/')
def index():
    # 返回HTML页面的内容
    return render_template_string(open('./Fronted/index.html').read())


if __name__ == "__main__":
    # app.run(debug=True)
    # 可选参数: host='0.0.0.0' 如果需要外部访问
    app.run(debug=True, host='0.0.0.0', port=8883)